{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Copyright (C) 2022-2023, Intel Corporation\n",
        "\n",
        "SPDX-License-Identifier: MIT\n",
        "\n",
        "# Quantization Aware Training using AzureML\n",
        "\n",
        "In this sample we are demonstrating quantization aware training using Neural Networks Compression Framework though Optimum Intel.   \n",
        "Here we're using a BERT Model from the HuggingFace hub (transformers library) for Question-Answering usecase.\n",
        "\n",
        "\n",
        "This notebook is made with the reference of examples mentioned on the below links:  \n",
        "- https://learn.microsoft.com/en-us/azure/machine-learning/how-to-train-with-custom-image\n",
        "- https://learn.microsoft.com/en-us/azure/machine-learning/quickstart-create-resources\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "# Getting Required Files\n",
        "\n",
        "We need to get the required files from the repository [here](https://github.com/intel/nlp-training-and-inference-openvino/tree/v1.1)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1670236174499
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "import os\n",
        "SCRIPT_DIR = 'bert_qat/'\n",
        "if not os.path.exists(SCRIPT_DIR):\n",
        "    os.mkdir(SCRIPT_DIR)\n",
        "if not os.path.exists(SCRIPT_DIR+\"trainer_qa.py\"):\n",
        "    !cd bert_qat/ && wget https://raw.githubusercontent.com/intel/nlp-training-and-inference-openvino/v1.1/question-answering-bert-qat/quantization_aware_training/training_scripts/trainer_qa.py\n",
        "\n",
        "if not os.path.exists(SCRIPT_DIR+\"run_qa.py\"):\n",
        "    !cd bert_qat/ && wget https://raw.githubusercontent.com/intel/nlp-training-and-inference-openvino/v1.1/question-answering-bert-qat/quantization_aware_training/training_scripts/run_qa.py\n",
        "\n",
        "if not os.path.exists(SCRIPT_DIR+\"utils_qa.py\"):\n",
        "    !cd bert_qat/ && wget https://raw.githubusercontent.com/intel/nlp-training-and-inference-openvino/v1.1/question-answering-bert-qat/quantization_aware_training/training_scripts/utils_qa.py\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "# Import necessary Libraries\n",
        "\n",
        "First we need to import the necessary libraries to perform the desired task"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "gather": {
          "logged": 1670236174855
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from pathlib import Path\n",
        "from azureml.core import Workspace\n",
        "from azureml.core import ScriptRunConfig, Experiment, Environment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "# Initialize Workspace\n",
        "\n",
        "The Azure Machine Learning workspace is the top-level resource for the service. It gives you a centralized place to work with all the artifacts that you create. In the Python SDK, you can access the workspace artifacts by creating a Workspace object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "gather": {
          "logged": 1670236177007
        }
      },
      "outputs": [],
      "source": [
        "from azureml.core import Workspace\n",
        "from azureml.core import Datastore\n",
        "ws = Workspace.from_config()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "# Create or attach a compute target\n",
        "\n",
        "A Compute target is a machine where we intend to run our code. It can be a compute instance or a compute clusters.  \n",
        "Here we are using a compute cluster. If the cluster already exists it'll attach our workspace to that cluster, else it'll create a cluster according to the specification mentioned and attach to our workspace."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1670236178132
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "from azureml.core.compute import ComputeTarget, AmlCompute\n",
        "from azureml.core.compute_target import ComputeTargetException\n",
        "\n",
        "# Choose a name for your cluster.\n",
        "cluster_name = \"cpu-cluster4\"\n",
        "\n",
        "try:\n",
        "    compute_target = ComputeTarget(workspace=ws, name=cluster_name)\n",
        "    print('Found existing compute target.')\n",
        "except ComputeTargetException:\n",
        "    print('Creating a new compute target...')\n",
        "    compute_config = AmlCompute.provisioning_configuration(vm_size='STANDARD_D16DS_V4',\n",
        "                                                           max_nodes=4)\n",
        "    # Create the cluster.\n",
        "    compute_target = ComputeTarget.create(ws, cluster_name, compute_config)\n",
        "\n",
        "    compute_target.wait_for_completion(show_output=True)\n",
        "\n",
        "# Use get_status() to get a detailed status for the current AmlCompute.\n",
        "print(compute_target.get_status().serialize())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "# Setting the correct input paths\n",
        "\n",
        "All scripts & files present in the `script_dir` script folder are uploaded to the compute target, data stores are mounted or copied, and the script is executed.  \n",
        "Outputs from stdout and the `./logs` folder are streamed to the run history and can be used to monitor the run. For further details please refer [here](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-train-pytorch#what-happens-during-run-execution)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "gather": {
          "logged": 1670236178538
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "script_dir = os.path.abspath(SCRIPT_DIR)+\"/\"\n",
        "script_name = \"run_qa.py\""
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "# Environment Definition\n",
        "\n",
        "Environment definition allows us to define a custom Docker Environment with all the required dependencies making sure our script runs as expected.\n",
        "\n",
        "## Location of the Dockerfile\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "gather": {
          "logged": 1670236178934
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "# Copyright (C) 2022-2023 Intel Corporation\n",
        "# SPDX-License-Identifier: MIT\n",
        "dockerfile = r\"\"\"\n",
        "FROM openvino/ubuntu20_runtime:2022.2.0\n",
        "\n",
        "USER root\n",
        "\n",
        "RUN dpkg --get-selections | grep -v deinstall | awk '{print $1}' > base_packages.txt\n",
        "\n",
        "RUN apt-get update && apt-get install -y wget\\\n",
        "    python3.8 \\\n",
        "    python3.8-venv; \\\n",
        "    rm -rf /var/lib/apt/lists/*;\n",
        "\n",
        "RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.8 70; \\\n",
        "    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.8 70;\n",
        "\n",
        "RUN apt-get update && apt-get install -y git\n",
        "\n",
        "RUN python -m pip install --upgrade pip\n",
        "\n",
        "RUN pip install --no-cache-dir \"git+https://github.com/huggingface/optimum-intel.git@v1.5.2#egg=optimum-intel[openvino,nncf]\"\n",
        "RUN pip install --no-cache-dir \"git+https://github.com/AlexKoff88/nncf_pytorch.git@ak/qdq_per_channel#egg=nncf\"\n",
        "RUN pip install --no-cache-dir \"protobuf==3.19.5\"\n",
        "RUN pip install --no-cache-dir \"seqeval==1.2.2\"\n",
        "RUN pip install --no-cache-dir \"accelerate==0.15.0\"\n",
        "RUN pip install --no-cache-dir \"evaluate==0.3.0\"\n",
        "RUN pip install --no-cache-dir \"datasets==2.7.1\"\n",
        "RUN pip install --no-cache-dir \"torch==1.12.0\"\n",
        "RUN pip install --no-cache-dir \"openvino-dev==2022.2.0\"\n",
        "WORKDIR /home/training\n",
        "\n",
        "RUN chown openvino -R /home/training\n",
        "\n",
        "USER openvino\n",
        "\n",
        "RUN ls -l\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Azure ML settings\n",
        "Assigning a name to our environment for easier tracking and monitoring"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "gather": {
          "logged": 1670236179188
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "environment_name = \"qat-example\"\n",
        "experiment_name = \"qat-test\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Creating Environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "gather": {
          "logged": 1670236179482
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "env = Environment(environment_name)\n",
        "env.docker.base_image = None\n",
        "env.docker.base_dockerfile = dockerfile\n",
        "env.python.user_managed_dependencies = True"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Arguments for Quantization Aware Training\n",
        "\n",
        "To run the Quantization Aware Training we need to pass few arguments to the script. The arguments can be easily passed by combining them together in a list in the below order:  \n",
        "`arguments = ['--first_arg', first_val, '--second_arg', second_val, ...]`\n",
        "\n",
        "For further details please refer [here](https://azure.github.io/azureml-cheatsheets/docs/cheatsheets/python/v1/cheatsheet/)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "gather": {
          "logged": 1670236179786
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "model_name =  \"bert-large-uncased-whole-word-masking-finetuned-squad\"\n",
        "arguments = [\"--model_name_or_path\", model_name,\n",
        "             \"--dataset_name\", \"squad\",\n",
        "             \"--do_train\", True,\n",
        "             \"--do_eval\", True,\n",
        "             \"--max_seq_length\", 256,\n",
        "             \"--per_device_train_batch_size\", 3,\n",
        "             \"--max_train_samples\",10,\n",
        "             \"--max_eval_samples\",10,\n",
        "             \"--learning_rate\", 3e-5,\n",
        "             \"--num_train_epochs\", 2,\n",
        "             \"--output_dir\", \"./outputs/bert_finetuned_model\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Create job config\n",
        "\n",
        "Job config allows us to define how we want to execute our training procedure. We need to pass the following informations to `ScriptRunConfig` object to initialize the job config instance.\n",
        "- `source_directory` $\\rightarrow$ All the contents of this source directory are copied to the compute target instance.  \n",
        "- `script` $\\rightarrow$ Our desired python script which we wish to execute  \n",
        "- `arguments` $\\rightarrow$ Necessary arguments for the `script`  \n",
        "- `env` $\\rightarrow$ Our target environment to execute the `script`\n",
        "- `compute_target` $\\rightarrow$ Our target compute preference (i.e. cluster or instance) to run the `script`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "gather": {
          "logged": 1670236180060
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "src = ScriptRunConfig(source_directory=script_dir,\n",
        "                      script=script_name,\n",
        "                      arguments=arguments,\n",
        "                      environment=env,\n",
        "                      compute_target=cluster_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "# Submit job\n",
        "After submitting the job, we can see logs from the Outputs + logs tab of the Web View link generated. Once job is completed, we can see output model from Web View link > Outputs + logs > outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1669972749717
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "run = Experiment(ws, experiment_name).submit(src)\n",
        "run.wait_for_completion(show_output=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "## Download the model outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "gather": {
          "logged": 1669888436945
        },
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [],
      "source": [
        "run.get_file_names()\n",
        "print('filenames',run.get_file_names())\n",
        "run.download_file(name='./outputs/bert_finetuned_model/model.onnx',\n",
        "                 output_file_path='./models/model.onnx')\n",
        "run.download_file(name='./outputs/bert_finetuned_model/config.json',\n",
        "                 output_file_path='./models/config.json')"
      ]
    }
  ],
  "metadata": {
    "kernel_info": {
      "name": "python38-azureml"
    },
    "kernelspec": {
      "display_name": "Python 3.8 - AzureML",
      "language": "python",
      "name": "python38-azureml"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.5"
    },
    "microsoft": {
      "host": {
        "AzureML": {
          "notebookHasBeenCompleted": true
        }
      }
    },
    "nteract": {
      "version": "nteract-front-end@1.0.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}
